import io
import os
import pickle
import zlib
from base64 import b64decode, b64encode
import SimpleITK as sitk
import numpy as np
# import pyAesCrypt
from PIL import Image
import pydicom



def to_npy(filepath, foldertype='Dicom',data_type="float32"):
    """cast ImageFolderType to numpy"""
    sitk_to_npy = DataTransformer.SitkToNpy
    # if foldertype == 'Nii':
    #     return sitk_to_npy(DataTransformer.NiiToSitk(filepath))
    if foldertype == 'Dicom':
        sitk_obj = DataTransformer.DicomToSitk(filepath)
        return sitk_obj, sitk_to_npy(sitk_obj)
    # elif foldertype == 'Mhd':
    #     return sitk_to_npy(DataTransformer.MhdToSitk(filepath,data_type))
    return None

def xy_spacing(dcm_files=None, all_image_tags=None):
    if not any([dcm_files, all_image_tags]):
        raise Exception('dcm_files and image_tags at least need one')
    xys = []
    if all_image_tags:
        pixel_spacing = all_image_tags[0].get('PixelSpacing')
        if pixel_spacing:
            xys = [float(x) for x in pixel_spacing]
    else:
        ds0 = pydicom.dcmread(dcm_files[0], force=True)
        if hasattr(ds0, "PixelSpacing"):
            xys = [float(x) for x in ds0.get("PixelSpacing")]
    if not xys:
        raise Exception('')
    return xys


def real_z_spacing(dcm_files=None, all_image_tags=None):
    cor = []
    if not any([dcm_files, all_image_tags]):
        raise Exception('dcm_files and image_tags at least need one')
    if all_image_tags:
        for image_tags in all_image_tags:
            slice_location = image_tags.get('SliceLocation')
            if slice_location:
                cor.append(float(slice_location))
            if len(cor) == 2:
                break
    else:
        for dcm_file in dcm_files:
            ds = pydicom.dcmread(dcm_file, force=True)
            if hasattr(ds, "SliceLocation"):
                cor.append(float(ds.get("SliceLocation")))
            if len(cor) == 2:
                break
    if len(cor) != 2:
        raise Exception("dicom file has no SliceLocation")
    z_spacing = np.linalg.norm(cor[1] - cor[0])
    print(f'will set z_spacing to {z_spacing}')
    return z_spacing

class DataTransformer:
    """
    转化数据类型
    DicomToSitk
    MhdToSitk
    SitkToMhd
    DcmToMhd
    NiiToSitk
    NiiToNpy
    NpyToSitk
    ------------------
    SitkToMhd
    DcmToMhd
    ------------------
    SitkToNii
    DicomToNii
    NumpyToNii
    ------------------
    Str2Code
    Code2Str
    Str2Pickle
    Pickle2Str
    Npy2Str
    Str2Npy
    Dcm2Np
    """

    @classmethod
    def DicomToSitk(cls, files):
        """
        SimpleITK unlocks the GIL, itk version can hang in multithreaded environments

        Args:
            files: list of dicom file names in the same series
        """
        reader = sitk.ImageSeriesReader()
        reader.SetFileNames(files)
        try:
            image = reader.Execute()
            xy = xy_spacing(files)
            image.SetSpacing((xy[0], xy[1], real_z_spacing(files)))
        except Exception as e:
            try:
                image = sitk.ReadImage(files)
            except:
                raise Exception("PixelData is invalid")
        return image

    @classmethod
    def MhdToSitk(cls, image_path,data_type="float32"):
        sitkImage = sitk.ReadImage(image_path)
        if data_type == "int16":
            sitkImage = sitk.Cast(sitkImage, sitk.sitkInt16)
        else:
            sitkImage = sitk.Cast(sitkImage, sitk.sitkFloat32)
        return sitkImage

    @classmethod
    def SitkToMhd(cls, sitk_img:sitk.Image, dest:str):
        assert dest.endswith(".mhd")
        sitk.WriteImage(sitk_img, dest)
        return dest


    @classmethod
    def DcmToMhd(cls, files:list, dest:str):
        assert dest.endswith(".mhd")
        sitk_img = cls.DicomToSitk(files)
        return cls.SitkToMhd(sitk_img, dest)


    @classmethod
    def NiiToSitk(cls, filepath):
        if isinstance(filepath, list):
            assert len(filepath) == 1, "data structure not understood, do cross-check with bmutils.generate_series"
            filepath = filepath[0]  # to accomodate bmutils.series.generate_series
        assert isinstance(filepath, str), "input should be file_path of nii.gz or nii"
        return sitk.ReadImage(filepath)

    
    @classmethod
    def NiiToNpy(cls, filepath):
        sitk_img = sitk.ReadImage(filepath)
        img = sitk.GetArrayFromImage(sitk_img)
        return img

    @classmethod
    def NpyToSitk(cls, array, sitk_ref, dtype="int32"):
        assert isinstance(array, np.ndarray)
        assert isinstance(sitk_ref, sitk.SimpleITK.Image)
        array = array.astype(dtype)
        image = sitk.GetImageFromArray(array)
        image.CopyInformation(sitk_ref)
        return image


    @classmethod
    def SitkToNii(cls, itk_img: sitk.Image, dest:str):
        assert dest.endswith(".nii") or dest.endswith(".nii.gz")
        return  sitk.WriteImage(itk_img, dest)

    @classmethod
    def DcmToNii(cls, files, dest):
        assert dest.endswith(".nii") or dest.endswith(".nii.gz")
        sitk_img = cls.DicomToSitk(files)
        return  cls.StikToNii(sitk_img, dest)

    @classmethod
    def NumpyToNii(cls, ndarray:np.ndarray, dest:str):
        assert dest.endswith(".nii") or dest.endswith(".nii.gz")
        sitk_img = cls.NpyToSitk(ndarray)
        return  cls.StikToNii(sitk_img, dest)


    @classmethod
    def SitkToNpy(cls, sitk_image, view=False):
        """
        SimpleITK unlocks the GIL, itk version can hang in multithreaded environments
        """
        if view:
            return sitk.GetArrayViewFromImage(sitk_image)
        else:
            return sitk.GetArrayFromImage(sitk_image)

    @classmethod
    def Npy2Str(cls, ndarr):
        return b64encode(zlib.compress(ndarr.dumps(), 9)).decode()

    @classmethod
    def Npy2StrV2(cls, ndarr):
        """
        Convert a numpy to base64 string
        """
        try:
            base64_mask = b64encode(ndarr).decode("utf-8")
        except Exception as e:
            logger.error("npy_to_str:[{0}]".format(e))
            return ""
        return base64_mask

    @classmethod
    def Str2Npy(cls, base6s_img, png=None):
        bytes_img = b64decode(base6s_img)
        if png or b"PNG" in bytes_img[:10]:
            img = Image.open(io.BytesIO(bytes_img))
            img = img.convert("1")
            mask_np = np.array(img, dtype=np.uint8)
        else:
            mask_np = pickle.loads(zlib.decompress(bytes_img))
        return mask_np

    @classmethod
    def Str2NpyL(cls, base64_img):
        bytes_img = b64decode(base64_img)
        img = Image.open(io.BytesIO(bytes_img))
        img = img.convert("L")
        mask_np = np.array(img, dtype=np.uint8)
        return mask_np

    @classmethod
    def Pickle2Str(cls, obj, level=4):
        return zlib.compress(pickle.dumps(obj), level)

    @classmethod
    def Str2Pickle(cls, s):
        return pickle.loads(zlib.decompress(s))

    # @classmethod
    # def Code2Str(cls, code, aes_key=None):
    #     code_b64 = b64encode(zlib.compress(code.encode(), 9)).decode()
    #     if aes_key is not None:
    #         BUFFER_SIZE = 64 * 1024
    #         ins = io.BytesIO()
    #         ins.write(code_b64.encode())
    #         ins.seek(0, os.SEEK_SET)
    #         outs = io.BytesIO()
    #         pyAesCrypt.encryptStream(ins, outs, aes_key, BUFFER_SIZE)
    #         outs.seek(0)
    #         code_b64_enc = outs.read()
    #         return code_b64_enc.hex()
    #     else:
    #         return code_b64

    # @classmethod
    # def Str2Code(cls, codestr, aes_key=None):
    #     if aes_key is not None:
    #         BUFFER_SIZE = 64 * 1024
    #         ins = io.BytesIO()
    #         ins.write(bytes.fromhex(codestr))
    #         ins.seek(0, os.SEEK_SET)
    #         outs = io.BytesIO()
    #         pyAesCrypt.decryptStream(ins, outs, aes_key, BUFFER_SIZE, len(ins.getvalue()))
    #         outs.seek(0)
    #         return zlib.decompress(b64decode(outs.read()))
    #     else:
    #         return zlib.decompress(b64decode(codestr))

    @classmethod
    def write_sitk_data_with_metadata(cls, file_path, mask_data):
        """
            input : filepath , data(np.ndarray or sitk) 
            output: numpy.array
        """
        try:
            """
                如果是sitk ，则：
                    1.获取spacing
                    2.转numpy ，设置 dtype为 uint8
                    3.转回sitk
                    4.设置spacing
                如果不是sitk，则为np.array：
                    1.转为sitk
            """
            if type(mask_data) is sitk.SimpleITK.Image:
                data_spacing = mask_data.GetSpacing()
                data_origin = mask_data.GetOrigin()
                data_direction = mask_data.GetDirection()
                source_np = cls.SitkToNpy(mask_data)
                data = sitk.GetImageFromArray(source_np.astype(np.uint8))
                data.SetSpacing(data_spacing)
                data.SetOrigin(data_origin)
                data.SetDirection(data_direction)
    
            elif type(mask_data) is np.ndarray:
                data = sitk.GetImageFromArray(mask_data.astype(np.uint8))
            else:
                logger.error(f'This mask data type is {type(mask_data)}')
                return
            sitk.WriteImage(data, file_path)
        except Exception as e:
            raise Exception(f"Can't Write sitk_data:{e}")

    @classmethod
    def write_sitk_data_with_source_data(cls, file_path, source_data, img_info={}, np_dtype='int16'):
        """
            储存模型配准后的原图到 nii.gz 文件
            input : filepath , data(np.ndarray or sitk), img_info(origin/spacing/direction)
            output: numpy.array
        """
        try:
            """
                如果是sitk ，则：
                    1.获取spacing
                    2.转numpy ，设置 dtype为 int16
                    3.转回sitk
                    4.设置spacing
                如果不是sitk，则为np.array：
                    1.转为sitk
                    2.设置spacing
            """
            if type(source_data) is sitk.SimpleITK.Image:
                data_spacing = source_data.GetSpacing()
                data_origin = source_data.GetOrigin()
                data_direction = source_data.GetDirection()
                source_np = cls.SitkToNpy(source_data)
                data = sitk.GetImageFromArray(source_np.astype(np.int16))

            elif type(source_data) is np.ndarray:
                if np_dtype == "int16":
                    data = sitk.GetImageFromArray(source_data.astype(np.int16))
                elif np_dtype == 'uint8':
                    data = sitk.GetImageFromArray(source_data.astype(np.uint8))
                data_spacing = img_info.get('spacing', [1, 1, 1])
                data_origin = img_info.get('origin', [0, 0, 0])
                data_direction = img_info.get('direction', [1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0])
            else:
                logger.error(f'This source data type is {type(source_data)}')
                return
            data.SetSpacing(data_spacing)
            data.SetOrigin(data_origin)
            data.SetDirection(data_direction)
            sitk.WriteImage(data, file_path)
        except Exception as e:
            raise Exception(f"Can't Write sitk_data:{e}")

    @classmethod
    def read_sitk_data(cls, filepath):
        """
            input : filepath
            output: numpy.array
        """
        try:
            itk_img = sitk.ReadImage(filepath)
            return sitk.GetArrayFromImage(itk_img)
        except Exception as e:
            raise Exception(f"Can't Read sitk_data:{e}")

    @classmethod
    def Dcm2Np(cls, filepath):
        data_sitk = sitk.ReadImage(filepath)
        data_sitk = sitk.GetArrayFromImage(data_sitk).astype(np.float32)
        return data_sitk